import pandas as pd
from globleConfig.database import Database
import warnings

"""
该类实现了sql语句
"""

class pgsql(object):

    def __init__(self):
        db = Database()
        warnings.filterwarnings('ignore')
        self.cnn = db.getcnn() #数据库连接
        with self.cnn:
            self.cur = self.cnn.cursor()  #数据库游标

    #将df所示日线信息插入到数据库中
    def insert_index_daily(self, df):
        insert_statemant = 'INSERT INTO index_daily (ts_code, trade_date, close, open, high, low, pre_close, change, pct_chg, vol, amount, turnover_rate) VALUES'
        if df.shape[0] == 0: return
        for index, row in df.iterrows():
            value = f"('{row['ts_code']}', '{row['trade_date']}', {row['close']}, {row['open']}, {row['high']}, {row['low']}, {row['pre_close']}, {row['change']}, {row['pct_chg']}, {row['vol']}, {row['amount']}, {row['turnover_rate']})"
            insert_statemant += value + ','
        #删除最后一个逗号
        insert_statemant = insert_statemant[:-1]
       # print(insert_statemant)
        self.cur.execute(insert_statemant)
        self.cnn.commit()

    # 从数据库中获取日线信息
    def get_index_daily(self, ts_code, start, end):
        sql_query = f"""
                        SELECT 
                            trade_date, close, open, high, low, pre_close, change, pct_chg, vol, amount, turnover_rate
                        FROM 
                            index_daily
                        WHERE 
                            ts_code = '{ts_code}' AND 
                            trade_date BETWEEN '{start}' AND '{end}'
                        ORDER BY 
                            trade_date;
                        """
        df = pd.read_sql_query(sql_query, self.cnn)
        return df

    # 获取ts_code 股票在数据库中有多少条日线数据
    def getCountOfdaily(self, ts_code):
        sql_query = f"""
                     SELECT 
                     COUNT(*)
                     FROM index_daily
                     WHERE ts_code = '{ts_code}'
                    """
        df = pd.read_sql_query(sql_query, self.cnn)
        return df

    # 获取某值股票最后k天的日线数据
    def getLastkOfdaily(self, ts_code, k):
        sql_query = f"""
                    WITH latest_trades AS (
                    SELECT *,
                           ROW_NUMBER() OVER (PARTITION BY ts_code ORDER BY trade_date DESC) as row_num
                    FROM index_daily
                    WHERE ts_code = '{ts_code}'
                    )
                    SELECT *
                    FROM latest_trades
                    WHERE row_num <= '{k}';
                     """
        return pd.read_sql_query(sql_query, self.cnn)

    # 获取某只股票的所有日线数据
    def get_all_index_daily(self, ts_code):
        #print(ts_code)
        sql_query = f"""
                        SELECT 
                            trade_date, close, open, high, low, pre_close, change, pct_chg, vol, amount, turnover_rate
                        FROM 
                            index_daily
                        WHERE ts_code = '{ts_code}'
                        """
        df = pd.read_sql_query(sql_query, self.cnn)
        return df

    # 从数据库中获取最大的日期
    def get_maxdate_daily(self, ts_code):
        sql_query = f"""
                    SELECT 
                        MAX(trade_date) AS latest_date
                    FROM 
                        index_daily
                    WHERE 
                        ts_code = '{ts_code}';
                    """
        df = pd.read_sql_query(sql_query, self.cnn)
        return df

    # 从数据库中获取最大的日期
    def get_maxdate_cashflow(self, ts_code):
        sql_query = f"""
                            SELECT 
                                MAX(ann_date) AS latest_date
                            FROM 
                                cashflow
                            WHERE 
                                ts_code = '{ts_code}';
                            """
        df = pd.read_sql(sql_query, self.cnn)
        return df


    # 从数据库中获取最大的日期
    def get_maxdate_balancesheet(self, ts_code):
        sql_query = f"""
                            SELECT 
                                MAX(ann_date) AS latest_date
                            FROM 
                                balancesheet
                            WHERE 
                                ts_code = '{ts_code}';
                            """
        df = pd.read_sql(sql_query, self.cnn)
        return df

    # 从数据库中获取最大的日期
    def get_maxdate_income(self, ts_code):
        sql_query = f"""
                            SELECT 
                                MAX(ann_date) AS latest_date
                            FROM 
                                income
                            WHERE 
                                ts_code = '{ts_code}';
                            """
        df = pd.read_sql(sql_query, self.cnn)
        return df

    # 插入turnover_rate列到日线数据中
    # df列包括'rutnover_rate', 'ts_code', 'trade_date'
    def insertoverturn_rate(self, df):
        sql_query = """
                            UPDATE 
                                index_daily
                            SET 
                                turnover_rate = %s
                            WHERE 
                                ts_code = %s and trade_date = %s
                     """
        data_to_update = [(row['turnover_rate'], row['ts_code'], row['trade_date']) for _, row in df.iterrows()]
        self.cur.executemany(sql_query, data_to_update)
        self.cnn.commit()

    # 将数据通过df插入数据库的table表中
    def insert_data_bydf(self, df, table):
        columns = df.columns.tolist()
        insert_statement = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({', '.join(['%s'] * len(columns))})"
        for _, row in df.iterrows():
            #print(_, row)
            self.cur.execute(insert_statement, list(row))
        self.cnn.commit()